#include "RSA.h"

#include <ctime>
#include <stdexcept>
#include <utility>

#include <gmp.h>

#include "../DEBUG_PRINT_MACRO.h"

#include "data_conversion.h"

#ifdef DEBUG_PRINTS_LEVEL_1
#include <iostream>
#endif

// 引入 Tracy 头文件
#include <Tracy.hpp>


template<int key_length>
void RSA<key_length>::generateKeyPair(RSA_privatekey& private_key, RSA_publickey& public_key, RSA_publicparameter& publicParameters) {
    // Tracy 范围插桩，用于测量整个 generateKeyPair 函数的性能
    ZoneScopedN("RSA::generateKeyPair");

#ifdef DEBUG_PRINTS_LEVEL_1
    printf("[RSA][info] generateKeyPair function called\n");
#endif
    // 选择密钥长度
    int key_size = key_length;

    // 初始化临时变量
    mpz_t p, q, n, phi_n, p_minus_1, q_minus_1, e, d, gcd_result;
    
    // Tracy 范围插桩，用于测量 mpz_inits 的性能
    {
        ZoneScopedN("RSA::generateKeyPair::mpz_inits");
        mpz_inits(p, q, n, phi_n, p_minus_1, q_minus_1, e, d, gcd_result, nullptr);
    }

    // 初始化随机数生成器
    gmp_randstate_t state;
    // Tracy 范围插桩，用于测量 gmp_randinit_default 的性能
    {
        ZoneScopedN("RSA::generateKeyPair::gmp_randinit_default");
        gmp_randinit_default(state);
    }

    // 设置随机数种子
    unsigned long seed = static_cast<unsigned long>(time(NULL));
    // Tracy 范围插桩，用于测量 gmp_randseed_ui 的性能
    {
        ZoneScopedN("RSA::generateKeyPair::gmp_randseed_ui");
        gmp_randseed_ui(state, seed);
    }

    // 最大尝试次数
    constexpr int MAX_ATTEMPTS = 1000;
    int attempts = 0;
    bool success = false;


    while (attempts++ < MAX_ATTEMPTS) {
        // Tracy 范围插桩，用于测量每次循环迭代的性能
        ZoneScopedN("RSA::generateKeyPair::Key Generation Attempt Loop");
        // --------------------------------------------
        // 1. 生成两个大素数 p 和 q，在不满足条件时重新生成
        // --------------------------------------------
        {
            ZoneScopedN("RSA::generateKeyPair::Generate Primes p and q");
            do {
                mpz_urandomb(p, state, key_size / 2); // 素数长度为 key_size / 2, 以保证 n 满足长度要求
                mpz_nextprime(p, p);
            } while (mpz_sizeinbase(p, 2) != key_size / 2);

            do {
                mpz_urandomb(q, state, key_size / 2);
                mpz_nextprime(q, q);
            } while (mpz_sizeinbase(q, 2) != key_size / 2);
        }
#ifdef DEBUG_PRINTS_LEVEL_3
        std::cout << "[RSA_generateKeyPair][debug]Generated prime p: "; mpz_out_str(stdout, 10, p); std::cout << std::endl;
        std::cout << "[RSA_generateKeyPair][debug]Generated prime q: "; mpz_out_str(stdout, 10, q); std::cout << std::endl;
#endif
        // -----------------------------
        // 2. 计算模数 n 和欧拉函数 phi(n)
        // -----------------------------
        {
            ZoneScopedN("RSA::generateKeyPair::Calculate n and phi_n");
            mpz_mul(n, p, q); // n = pq
            mpz_sub_ui(p_minus_1, p, 1);   // p - 1
            mpz_sub_ui(q_minus_1, q, 1);   // q - 1
            mpz_mul(phi_n, p_minus_1, q_minus_1); // phi_n = (p-1)(q-1)
        }
#ifdef DEBUG_PRINTS_LEVEL_3
        std::cout << "[RSA_generateKeyPair][debug]Generated n: "; mpz_out_str(stdout, 10, n); std::cout << std::endl;
        std::cout << "[RSA_generateKeyPair][debug]Generated phi_n: "; mpz_out_str(stdout, 10, phi_n); std::cout << std::endl;
#endif
        //---------------------------------------------------------
        // 3. 选择公钥指数 e，使得 gcd(e, phi_n) = 1 且 1 < e < phi_n
        //---------------------------------------------------------
        {
            ZoneScopedN("RSA::generateKeyPair::Choose Public Exponent e");
            mpz_set_ui(e, 65537); // 常用的公钥指数

            // 检查 gcd(e, phi_n) 是否为 1
            mpz_gcd(gcd_result, e, phi_n);
            if (mpz_cmp_ui(gcd_result, 1) != 0) {
#ifdef DEBUG_PRINTS_LEVEL_3
                std::cerr << "[RSA key generate][Error] gcd(e, phi_n) is not 1. try again with another e value." << std::endl;
#endif
                continue; // 失败：gcd(e, phi_n) != 1 ，重新开始
            }
        }
        //-----------------------------------------
        // 4. 计算私钥指数 d，使得 ed ≡ 1 (mod phi_n)
        //-----------------------------------------
        {
            ZoneScopedN("RSA::generateKeyPair::Calculate Private Exponent d");
            if (mpz_invert(d, e, phi_n) == 0) {
                continue; // 失败：不存在模逆，重新开始
            }
        }

        break; // 成功，退出循环
    }

    // 尝试次数过多，抛出异常
    if (attempts >= MAX_ATTEMPTS) {
        // Tracy 范围插桩，用于测量清理资源和抛出异常的性能
        {
            ZoneScopedN("RSA::generateKeyPair::Cleanup and Throw Error");
            mpz_clears(p, q, n, phi_n, p_minus_1, q_minus_1, e, d, gcd_result, nullptr);
            gmp_randclear(state);
        }
#ifdef DEBUG_PRINTS_LEVEL_3
        std::cerr << "[RSA key generate][Error] Exceeded maximum number of key generation attempts." << std::endl;
#endif
        throw std::runtime_error("Exceeded maximum number of key generation attempts.");
    }
    //-------------
    // 5. 返回密钥对
    //-------------
    {
        ZoneScopedN("RSA::generateKeyPair::Return Key Pair");
        private_key = RSA_privatekey(d, n);
        public_key = RSA_publickey(e, n);
        publicParameters = RSA_publicparameter{};
    }

    // 清理资源
    // Tracy 范围插桩，用于测量清理资源的性能
    {
        ZoneScopedN("RSA::generateKeyPair::Cleanup Resources");
        mpz_clears(p, q, n, phi_n, p_minus_1, q_minus_1, e, d, gcd_result, nullptr);
        gmp_randclear(state);
    }
}

template<int key_length>
RSA_signature RSA<key_length>::sign(const PlainMessage& message, const RSA_privatekey& private_key, const RSA_publicparameter& publicParameters) {
    // Tracy 范围插桩，用于测量整个 sign 函数的性能
    ZoneScopedN("RSA::sign");

#ifdef DEBUG_PRINTS_LEVEL_1
    printf("[RSA sign][info] sign function called\n");
#endif
    mpz_t hash_m, s;
    // Tracy 范围插桩，用于测量 mpz_init 的性能
    {
        ZoneScopedN("RSA::sign::mpz_init_sign");
        mpz_init(hash_m);
        mpz_init(s);
    }
    // ------------------------------
    // 1. 对消息进行哈希运算并转换为整数
    // ------------------------------
    int ret;
    // Tracy 范围插桩，用于测量 hash_value_to_hex_string 和 mpz_set_str 的性能
    {
        ZoneScopedN("RSA::sign::Hash Message to Integer");
        std::string hex_string;
        hash_value_to_hex_string(message.get_hash_data(), hex_string);
        ret = mpz_set_str(hash_m, hex_string.c_str(), 16);
        if (ret != 0) {
            auto k = message.get_hash_data();
            throw std::invalid_argument("Invalid hash data format.");
        }
    }

#ifdef DEBUG_PRINTS_LEVEL_3
    // printf("[RSA sign][info] Message to be signed: %s\n", message.get_hash_data());
    printf("[RSA sign][info] m: ");
    mpz_out_str(stdout, 10, hash_m);
    printf("\n");
#endif

    // 确保 m < n，如果 m >= n，则需要对 m 进行模运算
    // Tracy 范围插桩，用于测量模运算的性能
    {
        ZoneScopedN("RSA::sign::Modulo Operation (m < n)");
        if (mpz_cmp(hash_m, private_key.n) >= 0) {
            mpz_mod(hash_m, hash_m, private_key.n);
        }
    }

    // ---------------------
    // 2. 使用私钥加密生成签名
    // ---------------------
    // Tracy 范围插桩，用于测量 mpz_powm 的性能
    {
        ZoneScopedN("RSA::sign::Generate Signature (mpz_powm)");
        mpz_powm(s, hash_m, private_key.d, private_key.n); // 计算签名值：s = H(m)^d mod n
    }

    // 清理临时变量
    // Tracy 范围插桩，用于测量清理变量的性能
    {
        ZoneScopedN("RSA::sign::Clear hash_m");
        mpz_clear(hash_m);
    }

    return RSA_signature{ std::move(s) };
}

template<int key_length>
bool RSA<key_length>::verify(const PlainMessage& message, const RSA_publickey& public_key, const RSA_publicparameter& publicParameters, const RSA_signature& signature) {
    // Tracy 范围插桩，用于测量整个 verify 函数的性能
    ZoneScopedN("RSA::verify");

#ifdef DEBUG_PRINTS_LEVEL_1
    printf("[RSA sign][info] verify function called\n");
#endif
    mpz_t s, m, hash_m;
    // Tracy 范围插桩，用于测量 mpz_inits 的性能
    {
        ZoneScopedN("RSA::verify::mpz_inits_verify");
        mpz_inits(s, m, hash_m, nullptr);
    }

    // ------------------------------
    // 1. 对消息进行哈希运算并转换为整数
    // ------------------------------
    int ret;
    // Tracy 范围插桩，用于测量 hash_value_to_hex_string 和 mpz_set_str 的性能
    {
        ZoneScopedN("RSA::verify::Hash Message to Integer (Verify)");
        std::string hex_string;
        hash_value_to_hex_string(message.get_hash_data(), hex_string);
        ret = mpz_set_str(hash_m, hex_string.c_str(), 16); // 假设消息是十进制字符串
        if (ret != 0) {
            throw std::invalid_argument("Invalid hash data format.");
        }
    }

    // 确保 m < n，如果 m >= n，则需要对 m 进行模运算
    // Tracy 范围插桩，用于测量模运算的性能
    {
        ZoneScopedN("RSA::verify::Modulo Operation (m < n) Verify");
        if (mpz_cmp(hash_m, public_key.n) >= 0) {
            mpz_mod(hash_m, hash_m, public_key.n);
        }
    }

    // ------------------
    // 2. 使用公钥解密签名
    // ------------------
    // Tracy 范围插桩，用于测量 mpz_powm 的性能
    {
        ZoneScopedN("RSA::verify::Decrypt Signature (mpz_powm)");
        mpz_powm(m, signature.s, public_key.e, public_key.n); // m = s^e mod n
    }

    // 比较解密后的明文与原始哈希值
    // Tracy 范围插桩，用于测量比较和清理资源的性能
    {
        ZoneScopedN("RSA::verify::Compare and Cleanup");
        if (mpz_cmp(m, hash_m) == 0) {
            mpz_clears(s, m, hash_m, nullptr);
            return true; // 验证成功
        }
        else {
            mpz_clears(s, m, hash_m, nullptr);
            return false; // 验证失败
        }
    }
}

template class RSA<1024>;
template class RSA<2048>;
template class RSA<3072>;